import numpy as np
from torch.utils.data import Dataset, DataLoader


def load_data(path, is_train=True):
    """
    从数据集中加载句子和对应实体
        path: 数据集路径
        is_train: 指示是否为训练集
    """
    sentences = []
    entities = []
    with open(path, "r") as fp:
        lines = fp.readlines()
    if is_train:
        for i in range(0, len(lines), 3):
            sentences.append(lines[i].strip().lower())
            entry = list()
            for e in lines[i+1].strip().split("|"):
                if not e: continue
                _range, entity = e.split(" ")
                _range = [int(i) for i in _range.split(",")]
                entry.append([*(_range), entity])
            entities.append(entry)
    else:
        lines = [l.strip().lower() for l in lines]
        for i in range(0, len(lines)):
            sentences.append(lines[i])
    return sentences, entities

class EntityDictionary(object):
    """
    所有标签的字典，能够将标签转换为id，或将id转换为标签。
    """
    def __init__(self):
        self.i2e = None
        self.e2i = None
        self.entities = set()
    
    def populate(self, entities):
        """
        在给定实体集合上训练词典
        对每个实体Entity，会生成S#Entity、B#Entity、M#Entity、E#Entity四个不同的标签，除此以外还有O、PAD、UNK等
        """
        for entry in entities:
            for entity in entry:
                B, M, E, S = "B#"+entity[2], "M#"+entity[2], "E#"+entity[2], "S#"+entity[2]
                self.entities = self.entities.union({B, M, E, S})
        self.entities.add("O")
        self.entities = list(self.entities)
        i2e_tuple = zip(range(4, len(self.entities)+4), self.entities)
        e2i_tuple = zip(self.entities, range(4, len(self.entities)+4))
        self.i2e = dict(i2e_tuple)
        self.e2i = dict(e2i_tuple)
        # 添加PAD、UNK、BEG、END
        self.i2e[0] = "[PAD]"
        self.i2e[1] = "[UNK]"
        self.i2e[2] = "[BEG]"
        self.i2e[3] = "[END]"
        self.e2i["[PAD]"] = 0
        self.e2i["[UNK]"] = 1
        self.e2i["[BEG]"] = 2
        self.e2i["[END]"] = 3

    def __len__(self):
        return len(self.i2e)

    def __getitem__(self, key):
        if type(key)==str:
            return self.e2i[key]
        elif type(key)==int:
            return self.i2e[key]
        else:
            raise TypeError

class WordDictionary(object):
    """词的字典，能够将词转换为对应id，或将id转换为词"""
    def __init__(self):
        self.i2w = None
        self.w2i = None
        self.word_freq = dict()

    def populate(self, sentences):
        for sent in sentences:
            for word in sent.split(" "):
                self.word_freq[word] = self.word_freq.get(word, 0)+1
        words = list(self.word_freq.keys())
        i2w_tuple = zip(range(4, len(words)+4), words)
        w2i_tuple = zip(words, range(4, len(words)+4))
        self.i2w = dict(i2w_tuple)
        self.w2i = dict(w2i_tuple)

        self.i2w[0] = "[PAD]"
        self.i2w[1] = "[UNK]"
        self.i2w[2] = "[BEG]"
        self.i2w[3] = "[END]"
        self.w2i["[PAD]"] = 0
        self.w2i["[UNK]"] = 1
        self.w2i["[BEG]"] = 2
        self.w2i["[END]"] = 3

    def __len__(self):
        return len(self.i2w)

    def __getitem__(self, key):
        if isinstance(key, str):
            return self.w2i[key]
        elif isinstance(key, int):
            return self.i2w[key]
        else:
            raise TypeError

    def get(self, key, default):
        if isinstance(key, str):
            return self.w2i.get(key, default)
        elif isinstance(key, int):
            return self.i2w.get(key, default)
        else:
            raise TypeError
        

class NNERDataset(Dataset):
    """
    本次任务的训练集和验证集对应的数据集
    """
    def __init__(self, sentences, entities, d1:WordDictionary, d2: EntityDictionary):
        super(NNERDataset, self).__init__()
        self.data = []
        sentences = [s.strip().lower() for s in sentences]
        for i_data, sentence, entity in zip(range(len(sentences)), sentences, entities):
            word_ids = [d1[w] for w in sentence.split(" ")]
            word_ids.append(d1["[END]"])
            # 先什么都不管 ，直接顺序标注，基于的假设是没有任何一个word对应多个标签, 如果存在这种情况会直接覆盖掉
            default = d2["O"]
            label = [default for i in range(len(word_ids)-1)]
            for entry in entity:
                if entry[1]==entry[0]+1:
                    label[entry[0]] = d2["S#"+entry[2]]
                else:
                    label[entry[0]] = d2["B#"+entry[2]]
                    label[entry[1]-1] = d2["E#"+entry[2]]
                    for i in range(entry[0]+1, entry[1]-1):
                        label[i] = d2["M#"+entry[2]]
            label.append(d2["[END]"])
            self.data.append([word_ids, label])

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return len(self.data)


class NNERTestDataset(Dataset):
    """
    测试集对应的数据集结构
    """
    def __init__(self, sentences, d1:WordDictionary):
        super(NNERTestDataset, self).__init__()
        self.data = []
        sentences = [s.strip().lower() for s in sentences]
        # 只需要将sentence转换为id序列
        for i_data, sentence in enumerate(sentences):
            word_ids = [d1.get(w, d1["[UNK]"]) for w in sentence.split(" ")]
            word_ids.append(d1["[END]"])
            self.data.append(word_ids)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return len(self.data)


if __name__ == "__main__":
    train_sent, train_ent = load_data("./assets/train.txt")
    d2 = EntityDictionary()
    d2.populate(train_ent)

    d1 = WordDictionary()
    d1.populate(train_sent)

    dataset = NNERDataset(train_sent, train_ent, d1, d2)


